from typing import Literal

from langchain.chat_models import init_chat_model
from pydantic import BaseModel, Field

from .prompt import analyze_prompt


class Evolution(BaseModel):
    should_evolve: bool = Field(..., description=(
        "Indicates whether an evolution operation should be triggered for the new memory and its neighbors."
    ))
    actions: list[Literal["strengthen", "update_neighbor"]] = Field(..., description=(
        "List of actions to perform (max 2, no duplicates):\n"
        "- 'strengthen': establish a link between the new memory and a neighbor;\n"
        "- 'update_neighbor': update a neighbor’s summary or tags."
    ))
    suggested_connections: list[str] = Field(description=(
        "When 'strengthen' is in actions, lists the IDs of neighbor memories to connect to."
    ))
    tags_to_update: list[str] = Field(description=(
        "When 'strengthen' is in actions, the new memory’s tags to add or update; "
        "aligned with suggested_connections."
    ))
    new_summary_neighborhood: list[str] = Field(description=(
        "When 'update_neighbor' is in actions, provides the updated summaries for each neighbor "
        "that needs updating."
    ))
    new_tags_neighborhood: list[list[str]] = Field(description=(
        "When 'update_neighbor' is in actions, provides new tag lists for each neighbor; "
        "order corresponds to new_summary_neighborhood."
    ))


class LabelContent(BaseModel):
    keywords: list[str] = Field(..., description=(
        "several specific, distinct keywords that capture key concepts and terminology."
        "Order from most to least important."
        "Don't include keywords that are the name of the speaker or time."
        "At least three keywords, but don't be too redundant."
    ))
    summary: str = Field(..., description=(
        "one sentence summarizing: 1.Main topic/domain. 2.Key arguments/points. 3.Intended audience/purpose."
    ))
    tags: list[str] = Field(..., description=(
        "several broad categories/themes for classification."
        "Include domain, format, and type tags."
        "At least three tags, but don't be too redundant."
    ))


llm = init_chat_model(
    "deepseek:deepseek-chat"
)


evolve_llm = llm.with_structured_output(Evolution)


def analyze_content(content: str, prompt: str = analyze_prompt) -> BaseModel:
    analyze_llm = llm.with_structured_output(LabelContent)
    result = analyze_llm.invoke(prompt.format(content=content))
    return result
